#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File         :pos_control_offb.py
@Description  :位置环控制，主要用于领机的offb控制
@Time         :2024/05/30 16:54:00
@Author       :Lin Yuheng
@Version      :1.0
'''

import rospy
import tty, termios
import sys, select
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
from geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import time
import PyKDL
import math
import numpy as np

waypoint = [0, 0]
waypoint_start = [0, 0]

class Controller:
    def __init__(self):
        self.plane_id = 0
        self.vehicle_id = "plane_1"
        self.px = 0
        self.py = 0
        self.pz = 0
        self.vx, self.vy, self.vz = 0, 0, 0
        self.roll, self.pitch, self.yaw = 0, 0, 0
        self.phi_dot, self.theta_dot, self.psi_dot = 0, 0, 0
        
        self.vehicle_state = State()
        self.should_exit = False
        self.reach_start = False

        # subscribers
        self.local_pose_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/local_position/pose",
            PoseStamped,
            callback=self.local_pose_callback,
        )
        self.mavros_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/state",
            State,
            callback=self.mavros_state_callback,
        )

        # pub
        self.target_motion_pub = rospy.Publisher(
            self.vehicle_id + "/mavros/setpoint_position/local",
            PoseStamped,
            queue_size=10
        )

        self.target_raw_pose = PoseStamped()

    def local_pose_callback(self, msg: PoseStamped):
        self.px = msg.pose.position.x
        self.py = msg.pose.position.y
        self.pz = msg.pose.position.z

    # def local_vel_callback(self, msg: TwistStamped):
    #     self.vx = msg.twist.linear.x
    #     self.vy = msg.twist.linear.y
    #     self.vz = msg.twist.linear.z

    # def odom_callback(self, msg: Odometry):
    #     x = msg.pose.pose.orientation.x
    #     y = msg.pose.pose.orientation.y
    #     z = msg.pose.pose.orientation.z
    #     w = msg.pose.pose.orientation.w
    #     rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    #     self.roll = rot.GetRPY()[0]
    #     self.pitch = rot.GetRPY()[1]
    #     self.yaw = rot.GetRPY()[2]

    #     self.phi_dot = msg.twist.twist.angular.x
    #     self.theta_dot = msg.twist.twist.angular.y
    #     self.psi_dot = msg.twist.twist.angular.z

    def mavros_state_callback(self, msg: State):
        if self.vehicle_state.mode == State.MODE_PX4_OFFBOARD:
            if msg.mode != State.MODE_PX4_OFFBOARD:
                rospy.logwarn("Vehicle quit offboard mode, quitting!")
                self.should_exit = True
        self.vehicle_state = msg
        
    def interpolate_linear(self, p1, p2, t):
        return (p1[0] + (p2[0] - p1[0]) * t, p1[1] + (p2[1] - p1[1]) * t)

    def interpolate_arc(self, center, radius: float, start_angle: float, end_angle: float, t: float):
        start_angle = math.pi / 2 - start_angle
        end_angle = math.pi / 2 - end_angle
        angle = start_angle + (end_angle - start_angle) * t
        return (center[0] + radius * math.cos(angle), center[1] + radius * math.sin(angle))
    
    def distance_to_setpoint(self, setpoint):
        return math.sqrt((self.px - setpoint[0]) ** 2 + (self.py - setpoint[1]) ** 2)
    
    def calc_distance(self, point1, point2):
        return math.sqrt((point2[0] - point1[0]) ** 2 + (point1[1] - point2[1]) ** 2)

    def calc_circle_center(self, point1, point2):
        if point1[0] == point2[0]:
            return (point1[0], (point1[1]+point2[1])/2), abs(point1[1]-point2[1])/2
        elif point1[1] == point2[1]:
            return ((point1[0]+point2[0])/2, point1[1]), abs(point1[0]-point2[0])/2
        else:
            print("无法求圆心和半径")
            return(0,0)

    def start(self):
        print("start fixed-wing path control")
        rospy.init_node(
            self.vehicle_id + "_path_control_node"
        )
        rate = rospy.Rate(20) # 50hz发送好像有点过快
        
        # enu 坐标系
        upward = 80
        waypoint_start = [50.0, -100.0]
        speed_cruise = 16 # m/s
        
        waypoint = waypoint_start
        
        # 轨迹的关键点
        A = (50.0, 120.0)
        B = (250.0, 120.0)
        C = (250.0, -100.0)
        D = (50.0, -100.0)
        
        # 半圆的中心和半径
        # center1 = (100.0, 120)  # 半圆 1 的中心
        # radius1 = 100
        # center2 = (100.0, -120)  # 半圆 2 的中心
        # radius2 = 100
        center1, radius1 = self.calc_circle_center(A, B)
        center2, radius2 = self.calc_circle_center(C, D)
        print(center1, radius1)
        print(center2, radius2)
        
        # 分段长度
        len1 = self.calc_distance(D, A) # 第一段直线长度
        len2 = math.pi * radius1  # 第二段半圆长度
        len3 = self.calc_distance(B, C)  # 第三段直线长度
        len4 = math.pi * radius2  # 第四段半圆长度
        length_total = len1 + len2 + len3 + len4
        
        time_total = length_total / speed_cruise

        # 计算每段轨迹的参数范围
        t1 = len1 / length_total * time_total
        t2 = t1 + len2 / length_total * time_total
        t3 = t2 + len3 / length_total * time_total

        while not rospy.is_shutdown():
            current_time = rospy.Time.now()
            dist = self.distance_to_setpoint(waypoint)
            if self.should_exit:
                break
            if not self.reach_start:
                if dist < 30:
                    rospy.loginfo("Start path follow")
                    self.reach_start = True
                    t_init = current_time.to_sec()
            else:
                t = current_time.to_sec() - t_init
                
                while t >= time_total:
                    t = t - time_total
                if t <= t1:
                    # 第一段直线
                    waypoint = self.interpolate_linear(D, A, t / t1)
                elif t <= t2:
                    # 第二段半圆
                    waypoint = self.interpolate_arc(center1, radius1, -math.pi/2, math.pi/2, (t - t1) / (t2 - t1))
                elif t <= t3:
                    # 第三段直线
                    waypoint = self.interpolate_linear(B, C, (t - t2) / (t3 - t2))
                else:
                    # 第四段半圆
                    waypoint = self.interpolate_arc(center2, radius2, math.pi/2, 3 * math.pi/2, (t - t3) / (time_total - t3))

            self.target_raw_pose.pose.position.x = waypoint[0]
            self.target_raw_pose.pose.position.y = waypoint[1]
            self.target_raw_pose.pose.position.z = upward

            try:
                self.target_motion_pub.publish(self.target_raw_pose)
            except rospy.ROSException as e:
                pass
            rate.sleep()

            rospy.loginfo(
                "wp:(%f,%f); position: (%f, %f); distance:%f", waypoint[1], waypoint[0], self.py, self.px, dist
            )


if __name__ == "__main__":
    Con = Controller()
    try:
        Con.start()
    except rospy.ROSInterruptException:
        pass
